import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
plt.figure(figsize=(12, 6))
ax = plt.gca()
[ax.spines[i].set＿visible(False) for i in ['top', 'right']]


def gatt(m, t):
    '''甘特图
    m机器集
    t时间集
    '''
    for j in range(len(m)):  # 工序j
        i = m[j] - 1  # 机器编号i
        if j == 0:
            plt.barh(i, t[j])
            plt.text(np.sum(t[:j + 1]) / 8, i, 'J%sT%s' % ((j + 1), t[j]), color='white', size=8)
        else:
            plt.barh(i, t[j], left=(np.sum(t[:j])))
            plt.text(np.sum(t[:j]) + t[j] / 8, i, 'J%sT%s' % ((j + 1), t[j]), color='white', size=8)


def schedule_gantt(task_schedule:pd.DataFrame):
    ## task_type,node_id,plane_id,woker_id,start,end
    task_nums = len(task_schedule)
    workers_id = task_schedule.loc[:,'worker_id']

    workers = list(set(workers_id))
    workers.sort()
    worker_map = dict([ (id,i) for i,id in enumerate(workers)])
    colors = ['#F0F8FF','#FAEBD7','#00FFFF','#7FFFD4','#F0FFFF',\
              '#FFEBCD',  '#0000FF','#8A2BE2','#A52A2A','#DEB887','#5F9EA0','m']

    color_map = dict()

    for i in range(ord('A'),ord('M')):
        color_map[chr(i)] = colors[i-ord('A')]


    for idx in range(task_nums):
        task_type,worker_id,plane_id = task_schedule.loc[idx,\
                                ['task_type','worker_id','plane_id']]
        # worker_id = task_schedule.loc[idx,'worker_id']
        y = worker_map[worker_id]
        left = task_schedule.loc[idx,'start']
        width = task_schedule.loc[idx,'end'] - left
        plt.barh(y,width=width,left=left,color=color_map[task_type],edgecolor='r')
        plt.text(left+1/2, y,'%s,%s' % (plane_id,task_type), \
                 color='r', size=7)

    plt.yticks(np.arange(len(workers)), np.arange(1, len(workers) + 1))
    # plt.show()
    plt.savefig(r'.\static_play\out\gantt.png')

    return color_map





if __name__ == '__main__':
    '''测试代码'''
    m = np.random.randint(1, 7, 35)
    t = np.random.randint(15, 25, 35)
    gatt(m, t)
    plt.yticks(np.arange(max(m)), np.arange(1, max(m) + 1))
    plt.show()
